"""
@File    : rasterio_backend.py.py
@Time    : 2025/7/15 下午10:59
@Author  : xiashuobad
@Desc    : 

/**
 * _ooOoo_
 * o8888888o
 * 88" . "88
 * (| -_- |)
 *  O\ = /O
 * ___/`---'\____
 * .   ' \\| |// `.
 * / \\||| : |||// \
 * / _||||| -:- |||||- \
 * | | \\\ - /// | |
 * | \_| ''\---/'' | |
 * \ .-\__ `-` ___/-. /
 * ___`. .' /--.--\ `. . __
 * ."" '< `.___\_<|>_/___.' >'"".
 * | | : `- \`.;`\ _ /`;.`/ - ` : | |
 * \ \ `-. \_ __\ /__ _/ .-` / /
 * ======`-.____`-.___\_____/___.-`____.-'======
 * `=---=' bug泛滥 佛已瘫痪
"""
# rsio/core/backends/rasterio_backend.py
import numpy as np
import logging
from typing import Literal, Union
from pathlib import Path

from .gdal_backend import GDALImage
from ..utils import (type_check, require_python_library, PaletteConverter, ArrayShapeConverter, GeoConverter,
                     require_attr_not_empty, ImageModeError)

logger = logging.getLogger(__name__)


class RasterioImage(GDALImage, backend_name='rasterio'):
    open_mode = Literal['r', 'r+', 'w']

    @require_python_library('rasterio', 'rasterio')
    def _open(self):
        import rasterio
        if self.mode != 'w':
            self._img_data = rasterio.open(self.file_path, self.mode)

    @property
    @require_attr_not_empty('img_data')
    def img_array(self) -> np.ndarray:
        return self._img_data.read().squeeze()

    @img_array.setter
    @type_check
    def img_array(self, img_array: np.ndarray):
        from rasterio import MemoryFile
        img_array = ArrayShapeConverter.array_shape_to_rasterio(img_array)
        if not self.img_data:
            meta = {
                'driver': 'GTiff',
                'width': img_array.shape[2],  # 注意 rasterio 使用 (bands, height, width) 格式
                'height': img_array.shape[1],
                'count': img_array.shape[0],  # 波段数
                'dtype': img_array.dtype,
            }
            self._mem_file = MemoryFile()
            self._img_data = self._mem_file.open(**meta)
        self.img_data.write(img_array)

    def set_img_array(self, array: np.ndarray, ) -> 'GDALImage':
        self.img_array = array
        return self

    @staticmethod
    def colormap(dataset: 'rasterio.DatasetBase') -> dict[int, tuple[int, int, int, int]]:
        try:
            palette = dataset.colormap(1)
        except ValueError:
            palette = None

        return palette

    @property
    def img_data(self) -> 'rasterio.io.DatasetWriter':
        return self._img_data

    @img_data.setter
    def img_data(self, img_data: 'rasterio.DatasetReader | rasterio.io.DatasetWriter'):
        self._copy_data_to_memory(img_data)

    def set_img_data(self, img_data: 'rasterio.DatasetBase') -> 'RasterioImage':
        self.img_data = img_data
        return self

    def _copy_data_to_memory(self, img_data: 'rasterio.DatasetReader | rasterio.io.DatasetWriter'):
        import rasterio
        from rasterio import MemoryFile
        self._mem_file = MemoryFile()
        self._img_data: rasterio.io.DatasetWriter = self._mem_file.open(**img_data.profile)
        for _, window in img_data.block_windows():
            self._img_data.write(img_data.read(window=window), window=window)
        if palette := RasterioImage.colormap(img_data):
            self._img_data.write_colormap(1, palette)

    @GDALImage.require_write_mode
    @require_attr_not_empty('img_data')
    def save(self, save_path: str = None, **kwargs):
        if self.img_data.name.startswith('/vsimem') or (save_path and save_path != self.file_path):
            import rasterio
            save_path = save_path or self.file_path
            Path(save_path).parent.mkdir(parents=True, exist_ok=True)
            with rasterio.open(save_path, 'w', **self.img_data.profile) as dst:
                for _, window in self.img_data.block_windows():
                    dst.write(self.img_data.read(window=window), window=window)
                if palette := self.colormap(self.img_data):
                    dst.write_colormap(1, palette)

    @property
    @require_attr_not_empty('img_data')
    def shape(self) -> tuple[int, ...]:
        return super().shape

    @property
    @require_attr_not_empty('img_data')
    def size(self) -> tuple[int, int]:
        return self.width, self.height

    @property
    @require_attr_not_empty('img_data')
    def channel(self) -> int:
        return self.img_data.count

    @property
    @require_attr_not_empty('img_data')
    def width(self) -> int:
        return self.img_data.width

    @property
    @require_attr_not_empty('img_data')
    def height(self) -> int:
        return self.img_data.height

    @property
    @require_attr_not_empty('img_data')
    def dtype(self) -> str:
        return self._img_data.dtypes[0]

    @property
    @require_attr_not_empty('img_data')
    def geo_bounds(self) -> tuple[float, float, float, float]:
        left, bottom, right, top = self.img_data.bounds
        return left, top, right, bottom

    @property
    @require_attr_not_empty('img_data')
    def nodata(self) -> int | float:
        return self._img_data.nodata

    @nodata.setter
    @require_attr_not_empty('img_data')
    @type_check
    def nodata(self, nodata: int | float):
        self._img_data.nodata = nodata

    def set_nodata(self, nodata: int | float) -> 'RasterioImage':
        self.nodata = nodata
        return self

    @property
    @require_attr_not_empty('img_data')
    def geo_transform(self) -> 'affine.Affine':
        return self.img_data.transform

    @geo_transform.setter
    @require_attr_not_empty('img_data')
    @type_check
    def geo_transform(self, geo_transform: tuple[float, ...]):
        self._img_data.transform = GeoConverter.geo_transform2rasterio(geo_transform)

    def set_geo_transform(self, geo_transform: tuple[float, ...]) -> 'RasterioImage':
        self.geo_transform = geo_transform
        return self

    @property
    @require_attr_not_empty('img_data')
    def geo_projection(self) -> str:
        crs = self.img_data.crs
        if crs:
            return crs.wkt

    @geo_projection.setter
    @require_attr_not_empty('img_data')
    @type_check
    def geo_projection(self, geo_projection: str | int):
        if not geo_projection:
            logger.warning('[yellow]投影为空，设置无效！', extra={'markup': True})
        else:
            self.img_data.crs = GeoConverter.geo_projection2wkt(geo_projection)

    def set_geo_projection(self, geo_projection: str | int) -> 'RasterioImage':
        self.geo_projection = geo_projection
        return self

    def _checked_mode(self, mode: open_mode) -> open_mode:
        if mode not in self.open_mode.__args__:
            raise ImageModeError(f'不支持的打开模式：{mode}, 请使用以下模式：{self.open_mode.__args__}！')
        return mode

    @property
    @require_attr_not_empty('img_data')
    def palette(self) -> dict:
        return self.colormap(self.img_data)

    @palette.setter
    @require_attr_not_empty('img_data')
    @type_check
    def palette(self, palette: Union[list, tuple, dict, bytes, np.ndarray]):
        palette = PaletteConverter.convert('rasterio', palette)
        if self.channel == 1:
            self.img_data.write_colormap(1, palette)
        else:
            logger.warning('[yellow]当前图像为多波段，设置colormap无效！', extra={'markup': True})

    def set_palette(self, palette: Union[list, tuple, dict, bytes, np.ndarray]) -> 'RasterioImage':
        self.palette = palette
        return self

    @require_attr_not_empty('img_data')
    @type_check
    def resize(self,
               target_size: tuple[int, int] | int,
               keep_ratio: bool = False,
               pad: int = None,
               resample: int = None,
               box: list[int] | tuple[int, int, int, int] = None,
               **kwargs
               ) -> 'RasterioImage':
        # 处理目标尺寸
        if isinstance(target_size, int):
            target_size = (target_size, target_size)
        if target_size == self.size:
            return self
        target_width, target_height = target_size
        img_data = self.img_data
        if box:
            x_min, y_min, x_max, y_max = box
            img_data = self.crop(x_min, y_min, x_max - x_min, y_max - y_min).img_data
        # 处理保持比例的情况
        if keep_ratio:
            src_width, src_height, = img_data.width, img_data.height
            scale = min(target_width / src_width, target_height / src_height)
            new_width = int(src_width * scale)
            new_height = int(src_height * scale)
        else:
            new_width, new_height = target_width, target_height
        # 读取并重采样
        from rasterio.enums import Resampling
        resized_array = img_data.read(
            out_shape=(img_data.count, new_height, new_width),
            resampling=resample or Resampling.bilinear
        )

        # 更新元数据
        from affine import Affine
        from rasterio import MemoryFile
        self._mem_file = MemoryFile()
        profile = img_data.profile.copy()
        palette = self.colormap(img_data)
        profile.update({
            'width': new_width,
            'height': new_height,
            'transform': img_data.transform * Affine.scale(
                (img_data.width / new_width),
                (img_data.height / new_height)
            )
        })
        if pad is not None:
            # 计算在目标图像中的位置（居中）
            x_offset = (target_width - new_width) // 2
            y_offset = (target_height - new_height) // 2
            profile.update({
                'width': target_width,
                'height': target_height,
                'transform': img_data.transform * Affine.scale(img_data.width / new_width,
                                                               img_data.height / new_height
                                                               ) * Affine.translation(-x_offset, -y_offset)
            })
            temp_array = np.full((resized_array.shape[0], target_height, target_width), pad, dtype=resized_array.dtype)
            temp_array[:, y_offset:y_offset + new_height, x_offset:x_offset + new_width] = resized_array
            resized_array = temp_array
        img_data = self._mem_file.open(**profile)
        img_data.write(resized_array)
        new_image = self.__class__('', 'w')
        new_image._img_data = img_data
        palette and self.set_palette(palette)
        return new_image

    def reverse_rgb(self) -> 'RasterioImage':
        array_reversed = self.reversed_rgb(self.img_array, 'rasterio')
        self.img_data.write(array_reversed, self.channel == 1 or None)
        return self

    @require_attr_not_empty('img_data')
    @type_check
    def crop(self,
             x_off: int = 0,
             y_off: int = 0,
             window_x_size: int = None,
             window_y_size: int = None
             ) -> 'RasterioImage':
        from rasterio.windows import Window
        from rasterio import MemoryFile
        x_off, y_off, window_x_size, window_y_size = self._checked_crop_window(x_off, y_off, window_x_size,
                                                                               window_y_size)
        window = Window(x_off, y_off, window_x_size, window_y_size)
        cropped_array = self.img_data.read(window=window)
        self._mem_file = MemoryFile()
        profile = self.img_data.profile.copy()
        profile.update({
            'width': window_x_size,
            'height': window_y_size,
            'transform': self.img_data.window_transform(window)
        })
        palette = self.palette
        new_image = self.__class__('', 'w')
        img_data = self._mem_file.open(**profile)
        img_data.write(cropped_array)
        new_image._img_data = img_data
        palette and new_image.set_palette(palette)
        return new_image

    @require_attr_not_empty('img_data')
    @type_check
    def read_window_array(self, x_off: int = 0, y_off: int = 0, window_x_size: int = None, window_y_size: int = None,
                          band_indexes: int | list[int] = None, **kwargs) -> np.ndarray:
        if isinstance(band_indexes, int):
            band_indexes = [band_indexes]
        if isinstance(band_indexes, list):
            for i, band_index in enumerate(band_indexes):
                if band_index < 0 or band_index >= self.channel:
                    raise ValueError(f'波段索引：{band_indexes} 超出了范围：{tuple(range(self.channel))}！')
                band_indexes[i] += 1
        from rasterio.windows import Window
        return self.img_data.read(band_indexes, window=Window(x_off, y_off, window_x_size, window_y_size),
                                  **kwargs).squeeze()

    def show(self, title: str = None):
        super().show(title)

    def close(self):
        if hasattr(self, '_img_data'):
            self._img_data.close()
        if hasattr(self, '_mem_file'):
            self._mem_file.close()

    def reproject(self,
                  dst_srs_obj: 'int|str|dict|RasterioImage|rasterio.CRS',
                  dst_width=None,
                  dst_height=None,
                  num_threads: int = 2
                  ) -> 'RasterioImage':
        import rasterio
        from rasterio.warp import calculate_default_transform, reproject, Resampling
        from rasterio import MemoryFile

        if isinstance(dst_srs_obj, RasterioImage):
            dst_srs = dst_srs_obj.img_data.crs
            dst_width, dst_height = dst_srs_obj.size
        else:
            dst_srs = GeoConverter.geo_projection2wkt(dst_srs_obj)
        left, top, right, bottom = self.geo_bounds
        dst_transform, dst_width, dst_height = calculate_default_transform(
            self.img_data.crs, dst_srs, self.width, self.height, left, bottom, right, top, dst_width=dst_width,
            dst_height=dst_height
        )
        kwargs = self.img_data.meta.copy()
        kwargs.update({
            'crs': dst_srs,
            'transform': dst_transform,
            'width': dst_width,
            'height': dst_height
        })
        src_img_data = self.img_data
        self._mem_file = MemoryFile()
        self._img_data: rasterio.io.DatasetWriter = self._mem_file.open(**kwargs)
        for i in range(1, self.channel + 1):  # 处理每个波段
            reproject(
                source=rasterio.band(src_img_data, i),
                destination=rasterio.band(self.img_data, i),
                src_transform=src_img_data.transform,
                src_crs=src_img_data.crs,
                dst_transform=dst_transform,
                dst_crs=dst_srs,
                resampling=Resampling.bilinear,  # 重采样方法：可选项如nearest、cubic
                num_threads=num_threads
            )
        return self

    @require_attr_not_empty('img_array')
    def convert_to_rgb(self):
        new_image = RasterioImage(
            self.file_path.replace(self.img_suffix, f'_to_rgb{self.img_suffix}'),
            'w',
        )
        array_to_rgb = ArrayShapeConverter.array_shape_to_rgb(self.img_array)
        new_image.set_img_array(array_to_rgb).set_nodata(self.nodata).set_geo_transform(
            self.geo_transform).set_geo_projection(self.geo_projection)

        return new_image

    @staticmethod
    @require_python_library('rasterio', 'rasterio')
    @type_check
    def write_array(
            fp: str,
            array: np.ndarray,
            geo_transform: tuple[float, ...] = None,
            geo_projection: str = None,
            nodata: int | float = None,
            palette: Union[list, tuple, dict, bytes, np.ndarray] = None,
            **meta
    ):
        import rasterio
        array = ArrayShapeConverter.array_shape_to_rasterio(array)
        count = RasterioImage.array_channels(array)
        meta.update({
            'driver': 'GTiff',
            'width': array.shape[-1],
            'height': array.shape[-2],
            'count': count,
            'dtype': array.dtype,
            'nodata': nodata,
            'transform': GeoConverter.geo_transform2rasterio(geo_transform),
            'crs': geo_projection
        })
        with rasterio.open(fp, 'w', **meta) as dst:
            dst.write(array, count == 1 or None)
            if palette is not None:
                if count == 1:
                    dst.write_colormap(1, PaletteConverter.convert('rasterio', palette))
                else:
                    logger.warning(f'[yellow]图像波段数为：{count}，无法写入颜色表！仅单波段图像有效。',
                                   extra={'markup': True})
